package com.liguangting.agentsflex.agent;

import com.agentsflex.core.agent.LLMAgent;
import com.agentsflex.core.agent.Output;
import com.agentsflex.core.llm.Llm;
import com.agentsflex.core.message.AiMessage;
import com.agentsflex.llm.qwen.QwenLlm;
import com.agentsflex.llm.qwen.QwenLlmConfig;
import com.liguangting.agentsflex.constant.ApiKeys;

import java.util.HashMap;
import java.util.Map;


/**
 * @author liguangting
 * @date 2024/7/9
 * @description
 */
public class SqlAgent extends LLMAgent {

    public SqlAgent(Llm llm) {
        super(llm, "您现在是一个 MySQL 数据库架构师，请根据如下的表结构信息，" +
                "帮我生成可以执行的 DDL 语句，以方便我用于创建 MySQL 的表结构。\n" +
                "注意：\n" +
                "请直接返回 DDL 内容，不需要解释，不需要涉及除了 DDL 语句以外的其他内容。\n" +
                "\n以下是表信息的内容：\n\n{ddlInfo}");
    }

    @Override
    protected Output onMessage(AiMessage aiMessage) {
        String sqlContent = aiMessage.getContent()
                .replace("```sql", "")
                .replace("```", "");
        return Output.ofDefault(sqlContent);
    }

    public static void main(String[] args) {
        QwenLlmConfig config = new QwenLlmConfig();
        config.setApiKey(ApiKeys.QWEN_API_KEY);
        Llm llm = new QwenLlm(config);

        SqlAgent agent = new SqlAgent(llm);

        String ddlInfo = "表名 student，字段 id,name";

        Map<String, Object> variables = new HashMap<>();
        variables.put("ddlInfo", ddlInfo);

        Output output = agent.execute(variables, null);
        System.out.println(output.getValue());
    }
}
